#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import logging
import subprocess
import time

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class TestConfig:
    """测试配置类"""
    def __init__(self, use_gpu=True, auto_select_gpu=False, gpu_device_id=0):
        self.USE_GPU = use_gpu
        self.AUTO_SELECT_GPU = auto_select_gpu
        self.GPU_DEVICE_ID = gpu_device_id
        self.GPU_COUNT = 1
        self.GPU_MEMORY_THRESHOLD = 10
        # 添加模型路径配置以避免初始化时出错
        self.MODEL_PATH = '/mnt/ssd/jsj/models/mock_model'


def print_gpu_status(label=""):
    """打印当前GPU状态"""
    try:
        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
        if label:
            logger.info(f"{label}GPU状态:\n{result.stdout}")
        else:
            logger.info(f"当前GPU状态:\n{result.stdout}")
    except Exception as e:
        logger.error(f"获取GPU状态失败: {e}")


def print_process_gpu_usage(pid=None):
    """打印指定进程或当前进程的GPU使用情况"""
    try:
        # 如果未指定PID，使用当前进程ID
        target_pid = pid if pid is not None else os.getpid()
        logger.info(f"检查进程ID: {target_pid} 的GPU使用情况")
        
        # 使用nvidia-smi查询进程的GPU使用情况
        result = subprocess.run(
            ['nvidia-smi', '--query-compute-apps=pid,gpu_uuid,used_gpu_memory,gpu_name', '--format=csv,noheader,nounits'],
            capture_output=True, text=True
        )
        
        logger.info(f"进程(PID: {target_pid})的GPU使用情况:")
        logger.info("pid, gpu_uuid, used_gpu_memory [MiB], gpu_name")
        
        # 记录当前进程的GPU使用情况
        current_process_usage = []
        
        # 打印所有GPU上的进程信息
        for line in result.stdout.strip().split('\n'):
            if line.strip():
                logger.info(line)
                # 检查是否包含目标进程ID
                if str(target_pid) in line.split(',')[0].strip():
                    current_process_usage.append(line)
                    logger.info(f"发现目标进程在GPU上运行: {line}")
        
        if not current_process_usage:
            logger.warning(f"未找到进程 {target_pid} 的GPU使用记录")
            
        return current_process_usage
    except Exception as e:
        logger.error(f"获取进程GPU使用情况失败: {e}")
        return []


def test_environment_before_import():
    """测试在导入任何库之前设置环境变量"""
    logger.info("\n===== 测试1: 在导入任何库之前设置环境变量 =====")
    
    # 清除之前可能设置的环境变量
    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        del os.environ['CUDA_VISIBLE_DEVICES']
        logger.info("已清除之前的CUDA_VISIBLE_DEVICES设置")
    
    # 设置CUDA_VISIBLE_DEVICES
    gpu_id = 1
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    logger.info(f"已设置环境变量CUDA_VISIBLE_DEVICES={gpu_id}")
    logger.info(f"当前环境变量: CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
    
    # 此时不导入任何CUDA相关库，直接检查环境变量
    logger.info("测试1完成: 环境变量已设置，尚未导入任何库")


def test_torch_import_only():
    """测试只导入torch的情况"""
    logger.info("\n===== 测试2: 只导入torch的情况 =====")
    
    # 清除之前可能设置的环境变量
    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        del os.environ['CUDA_VISIBLE_DEVICES']
        logger.info("已清除之前的CUDA_VISIBLE_DEVICES设置")
    
    # 设置CUDA_VISIBLE_DEVICES
    gpu_id = 1
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    logger.info(f"已设置环境变量CUDA_VISIBLE_DEVICES={gpu_id}")
    
    # 打印测试前的GPU状态
    print_gpu_status("测试前的")
    
    # 导入torch
    logger.info("开始导入torch...")
    import torch
    logger.info(f"导入torch后环境变量: CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
    
    # 检查CUDA可用性
    cuda_available = torch.cuda.is_available()
    device_count = torch.cuda.device_count() if cuda_available else 0
    logger.info(f"torch.cuda.is_available(): {cuda_available}")
    logger.info(f"torch.cuda.device_count(): {device_count}")
    
    # 创建一个大型张量
    if cuda_available:
        logger.info("创建一个大型测试张量...")
        large_tensor = torch.randn(10000, 10000, device='cuda:0')
        logger.info(f"成功在设备 cuda:0 上创建张量")
        
        # 检查内存使用情况
        current_device = torch.cuda.current_device()
        allocated_memory = torch.cuda.memory_allocated(current_device) / 1024**3  # 转换为GB
        logger.info(f"设备 cuda:{current_device} 已分配内存: {allocated_memory:.2f} GB")
    
    # 打印测试后的GPU状态
    print_gpu_status("测试后的")
    
    # 检查进程GPU使用情况
    print_process_gpu_usage()
    
    logger.info("测试2完成: 只导入torch并创建张量")


def test_model_handler_import_sequence():
    """测试ModelHandler的导入顺序和环境变量设置"""
    logger.info("\n===== 测试3: ModelHandler的导入顺序和环境变量设置 =====")
    
    # 清除之前可能设置的环境变量
    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        del os.environ['CUDA_VISIBLE_DEVICES']
        logger.info("已清除之前的CUDA_VISIBLE_DEVICES设置")
    
    # 先创建配置对象
    gpu_id = 1
    config = TestConfig(use_gpu=True, auto_select_gpu=False, gpu_device_id=gpu_id)
    logger.info(f"已创建配置对象，指定GPU ID: {gpu_id}")
    
    # 打印测试前的GPU状态
    print_gpu_status("测试前的")
    
    # 导入ModelHandler并初始化
    logger.info("开始导入ModelHandler...")
    try:
        from model_handler import ModelHandler, setup_cuda_environment
        
        # 手动调用setup_cuda_environment
        logger.info("手动调用setup_cuda_environment函数...")
        selected_gpus = setup_cuda_environment(config)
        logger.info(f"setup_cuda_environment返回的选中GPU: {selected_gpus}")
        logger.info(f"setup_cuda_environment后环境变量: CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
        
        # 现在导入torch（如果还没有导入的话）
        if 'torch' not in sys.modules:
            logger.info("导入torch...")
            import torch
        
        # 初始化ModelHandler
        logger.info("初始化ModelHandler...")
        handler = ModelHandler(config)
        logger.info(f"ModelHandler初始化后环境变量: CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
        
        # 验证CUDA设备
        if torch.cuda.is_available():
            logger.info("验证CUDA设备...")
            # 创建一个小型张量测试
            test_tensor = torch.randn(100, 100, device='cuda:0')
            logger.info(f"成功在设备 cuda:0 上创建测试张量")
            
            # 检查内存使用情况
            current_device = torch.cuda.current_device()
            device_name = torch.cuda.get_device_name(current_device)
            allocated_memory = torch.cuda.memory_allocated(current_device) / 1024**3
            logger.info(f"设备 cuda:{current_device} ({device_name}) 已分配内存: {allocated_memory:.2f} GB")
        
    except Exception as e:
        logger.error(f"导入或初始化ModelHandler失败: {e}")
    
    # 打印测试后的GPU状态
    print_gpu_status("测试后的")
    
    # 检查进程GPU使用情况
    print_process_gpu_usage()
    
    logger.info("测试3完成: ModelHandler导入和初始化测试")


def test_transformers_import():
    """测试transformers库的导入对CUDA设备的影响"""
    logger.info("\n===== 测试4: transformers库的导入对CUDA设备的影响 =====")
    
    # 清除之前可能设置的环境变量
    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        del os.environ['CUDA_VISIBLE_DEVICES']
        logger.info("已清除之前的CUDA_VISIBLE_DEVICES设置")
    
    # 设置CUDA_VISIBLE_DEVICES
    gpu_id = 1
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    logger.info(f"已设置环境变量CUDA_VISIBLE_DEVICES={gpu_id}")
    
    # 打印测试前的GPU状态
    print_gpu_status("测试前的")
    
    # 先导入torch
    logger.info("导入torch...")
    import torch
    
    # 再导入transformers
    logger.info("导入transformers...")
    from transformers import AutoModelForCausalLM, AutoTokenizer
    logger.info(f"导入transformers后环境变量: CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
    
    # 尝试创建一个小型模型
    try:
        logger.info("创建一个小型transformers模型...")
        # 使用一个非常小的模型进行测试
        model = AutoModelForCausalLM.from_pretrained("gpt2", device_map="auto", torch_dtype=torch.float16)
        logger.info(f"成功创建transformers模型，设备映射: {model.hf_device_map}")
        
        # 检查内存使用情况
        if torch.cuda.is_available():
            current_device = torch.cuda.current_device()
            allocated_memory = torch.cuda.memory_allocated(current_device) / 1024**3
            logger.info(f"设备 cuda:{current_device} 已分配内存: {allocated_memory:.2f} GB")
    except Exception as e:
        logger.error(f"创建transformers模型失败: {e}")
    
    # 打印测试后的GPU状态
    print_gpu_status("测试后的")
    
    # 检查进程GPU使用情况
    print_process_gpu_usage()
    
    logger.info("测试4完成: transformers库导入和模型创建测试")


def main():
    """主函数，运行所有测试"""
    logger.info("===== 开始综合GPU测试 =====")
    
    # 1. 测试在导入任何库之前设置环境变量
    test_environment_before_import()
    time.sleep(2)  # 给系统一些时间
    
    # 2. 测试只导入torch的情况
    test_torch_import_only()
    time.sleep(2)  # 给系统一些时间
    
    # 3. 测试ModelHandler的导入顺序和环境变量设置
    test_model_handler_import_sequence()
    time.sleep(2)  # 给系统一些时间
    
    # 4. 测试transformers库的导入对CUDA设备的影响
    test_transformers_import()
    
    logger.info("\n===== 综合GPU测试完成 =====")
    logger.info("\n测试总结:")
    logger.info("1. 如果测试2成功但测试3失败，问题可能出在ModelHandler的实现上")
    logger.info("2. 如果测试4失败，问题可能与transformers库的设备映射机制有关")
    logger.info("3. 请检查日志输出，特别注意各测试阶段的CUDA_VISIBLE_DEVICES设置和GPU内存分配情况")


if __name__ == "__main__":
    main()